import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

train_dataset = torchvision.datasets.CIFAR10('../data/CIFAR10', train=True,
                                             transform=torchvision.transforms.ToTensor(),
                                             download=True)

dataloader = DataLoader(train_dataset, batch_size=64,  shuffle=False)


class MyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.maxpool = nn.MaxPool2d(3, 1, ceil_mode=True)

    def forward(self, x):
        return self.maxpool(x)


mynet = MyNet()

writer = SummaryWriter("../logs_maxpool")
i = 0
for data_, label_ in dataloader:
    print(data_.shape, label_)
    writer.add_images("input", data_, i)
    output = mynet(data_)
    writer.add_images("output", output, i)
    i += 1
writer.close()
